import torch

from torch import nn


class BasicConv2d(nn.Module):
	def __init__(self, in_channels, out_channels, kernel_size, stride, padding=0):
		super(BasicConv2d, self).__init__()
		self.conv = nn.Conv2d(
			in_channels=in_channels,
			out_channels=out_channels,
			kernel_size=kernel_size,
			stride=stride,
			padding=padding,
			bias=False)
		self.bn = nn.BatchNorm2d(
			num_features=out_channels,
			eps=0.001,
			momentum=0.1,
			affine=True)
		self.relu = nn.ReLU(inplace=False)

	def forward(self, x):
		x = self.conv(x)
		x = self.bn(x)
		x = self.relu(x)
		return x


class Stem(nn.Module):
	def __init__(self, in_channels, out_channels):
		super(Stem, self).__init__()
		self.conv2d_1 = BasicConv2d(in_channels, 32, kernel_size=3, stride=2)
		self.conv2d_2 = BasicConv2d(32, 32, kernel_size=3, stride=1)
		self.conv2d_3 = BasicConv2d(32, 64, kernel_size=3, stride=1, padding=1)
		self.max_pool_1 = nn.MaxPool2d(3, stride=2)
		self.conv2d_4 = BasicConv2d(64, 80, kernel_size=1, stride=1)
		self.conv2d_5 = BasicConv2d(80, 192, kernel_size=3, stride=1)
		self.conv2d_6 = BasicConv2d(192, out_channels, kernel_size=3, stride=2)

	def forward(self, x):
		x = self.conv2d_1(x)
		x = self.conv2d_2(x)
		x = self.conv2d_3(x)
		x = self.max_pool_1(x)
		x = self.conv2d_4(x)
		x = self.conv2d_5(x)
		x = self.conv2d_6(x)
		return x


class InceptionResNetA(nn.Module):
	def __init__(self, scale=1.0):
		super(InceptionResNetA, self).__init__()
		self.scale = scale
		self.branch0 = BasicConv2d(256, 32, kernel_size=1, stride=1)
		self.branch1 = nn.Sequential(
			BasicConv2d(256, 32, kernel_size=1, stride=1),
			BasicConv2d(32, 32, kernel_size=3, stride=1, padding=1))
		self.branch2 = nn.Sequential(
			BasicConv2d(256, 32, kernel_size=1, stride=1),
			BasicConv2d(32, 32, kernel_size=3, stride=1, padding=1),
			BasicConv2d(32, 32, kernel_size=3, stride=1, padding=1))
		self.conv2d = nn.Conv2d(96, 256, kernel_size=1, stride=1)
		self.relu = nn.ReLU(inplace=False)

	def forward(self, x):
		x0 = self.branch0(x)
		x1 = self.branch1(x)
		x2 = self.branch2(x)
		out = torch.cat((x0, x1, x2), 1)
		out = self.conv2d(out)
		out = out * self.scale + x
		out = self.relu(out)
		return out


class InceptionResNetB(nn.Module):
	def __init__(self, scale=1.0):
		super(InceptionResNetB, self).__init__()
		self.scale = scale
		self.branch0 = BasicConv2d(896, 128, kernel_size=1, stride=1)
		self.branch1 = nn.Sequential(
			BasicConv2d(896, 128, kernel_size=1, stride=1),
			BasicConv2d(128, 128, kernel_size=(1, 7), stride=1, padding=(0, 3)),
			BasicConv2d(128, 128, kernel_size=(7, 1), stride=1, padding=(3, 0)))
		self.conv2d = nn.Conv2d(256, 896, kernel_size=1, stride=1)
		self.relu = nn.ReLU(inplace=False)

	def forward(self, x):
		x0 = self.branch0(x)
		x1 = self.branch1(x)
		out = torch.cat((x0, x1), 1)
		out = self.conv2d(out)
		out = out * self.scale + x
		out = self.relu(out)
		return out


class InceptionResNetC(nn.Module):
	def __init__(self, scale=1.0, noReLU=False):
		super(InceptionResNetC, self).__init__()
		self.scale = scale
		self.noReLU = noReLU
		self.branch0 = BasicConv2d(1792, 192, kernel_size=1, stride=1)
		self.branch1 = nn.Sequential(
			BasicConv2d(1792, 192, kernel_size=1, stride=1),
			BasicConv2d(192, 192, kernel_size=(1, 3), stride=1, padding=(0, 1)),
			BasicConv2d(192, 192, kernel_size=(3, 1), stride=1, padding=(1, 0)))
		self.conv2d = nn.Conv2d(384, 1792, kernel_size=1, stride=1)
		if not self.noReLU:
			self.relu = nn.ReLU(inplace=False)

	def forward(self, x):
		x0 = self.branch0(x)
		x1 = self.branch1(x)
		out = torch.cat((x0, x1), 1)
		out = self.conv2d(out)
		out = out * self.scale + x
		if not self.noReLU:
			out = self.relu(out)
		return out


class ReductionA(nn.Module):
	def __init__(self):
		super(ReductionA, self).__init__()
		self.branch0 = BasicConv2d(256, 384, kernel_size=3, stride=2)
		self.branch1 = nn.Sequential(
			BasicConv2d(256, 192, kernel_size=1, stride=1),
			BasicConv2d(192, 192, kernel_size=3, stride=1, padding=1),
			BasicConv2d(192, 256, kernel_size=3, stride=2))
		self.branch2 = nn.MaxPool2d(3, stride=2)

	def forward(self, x):
		x0 = self.branch0(x)
		x1 = self.branch1(x)
		x2 = self.branch2(x)
		out = torch.cat((x0, x1, x2), 1)
		return out


class ReductionB(nn.Module):
	def __init__(self):
		super(ReductionB, self).__init__()
		self.branch0 = nn.Sequential(
			BasicConv2d(896, 256, kernel_size=1, stride=1),
			BasicConv2d(256, 384, kernel_size=3, stride=2))
		self.branch1 = nn.Sequential(
			BasicConv2d(896, 256, kernel_size=1, stride=1),
			BasicConv2d(256, 256, kernel_size=3, stride=2))
		self.branch2 = nn.Sequential(
			BasicConv2d(896, 256, kernel_size=1, stride=1),
			BasicConv2d(256, 256, kernel_size=3, stride=1, padding=1),
			BasicConv2d(256, 256, kernel_size=3, stride=2))
		self.branch3 = nn.MaxPool2d(3, stride=2)

	def forward(self, x):
		x0 = self.branch0(x)
		x1 = self.branch1(x)
		x2 = self.branch2(x)
		x3 = self.branch3(x)
		out = torch.cat((x0, x1, x2, x3), 1)
		return out


class InceptionResnetV1(nn.Module):
	def __init__(self):
		super(InceptionResnetV1, self).__init__()
		self.stem = Stem(in_channels=3, out_channels=256)
		self.inception_res_net_a_5 = nn.Sequential(
			InceptionResNetA(scale=0.17),
			InceptionResNetA(scale=0.17),
			InceptionResNetA(scale=0.17),
			InceptionResNetA(scale=0.17),
			InceptionResNetA(scale=0.17))
		self.reduction_a = ReductionA()
		self.inception_res_net_b_10 = nn.Sequential(
			InceptionResNetB(scale=0.10),
			InceptionResNetB(scale=0.10),
			InceptionResNetB(scale=0.10),
			InceptionResNetB(scale=0.10),
			InceptionResNetB(scale=0.10),
			InceptionResNetB(scale=0.10),
			InceptionResNetB(scale=0.10),
			InceptionResNetB(scale=0.10),
			InceptionResNetB(scale=0.10),
			InceptionResNetB(scale=0.10))
		self.reduction_b = ReductionB()
		self.inception_resNet_c_5 = nn.Sequential(
			InceptionResNetC(scale=0.20),
			InceptionResNetC(scale=0.20),
			InceptionResNetC(scale=0.20),
			InceptionResNetC(scale=0.20),
			InceptionResNetC(scale=0.20))
		self.block8 = InceptionResNetC(noReLU=True)
		self.average_pooling = nn.AdaptiveAvgPool2d(1)
		self.dropout = nn.Dropout(p=0.8)

	def forward(self, x):
		x = self.stem(x)
		x = self.inception_res_net_a_5(x)
		x = self.reduction_a(x)
		x = self.inception_res_net_b_10(x)
		x = self.reduction_b(x)
		x = self.inception_resNet_c_5(x)
		x = self.block8(x)
		x = self.average_pooling(x)
		x = self.dropout(x)
		return x
